package com.spring.dict.maven;

import com.google.auto.service.AutoService;
import com.sun.source.tree.Tree;
import com.sun.tools.javac.api.JavacTrees;
import com.sun.tools.javac.code.Flags;
import com.sun.tools.javac.code.Symbol;
import com.sun.tools.javac.processing.JavacProcessingEnvironment;
import com.sun.tools.javac.tree.JCTree;
import com.sun.tools.javac.tree.TreeMaker;
import com.sun.tools.javac.tree.TreeTranslator;
import com.sun.tools.javac.util.*;

import javax.annotation.processing.*;
import javax.lang.model.SourceVersion;
import javax.lang.model.element.Element;
import javax.lang.model.element.TypeElement;
import java.util.Set;

/**
 * @author liuzhiqiang
 */
@SupportedAnnotationTypes("com.spring.dict.maven.Describe")
@SupportedSourceVersion(value = SourceVersion.RELEASE_8)
@AutoService(Processor.class)
public class DescribeAnnotationProcessor extends AbstractProcessor {

    /**
     * 将Element转换为JCTree的工具,提供了待处理的抽象语法树
     */
    private JavacTrees trees;

    /**
     * 封装了创建AST节点的一些方法
     */
    private TreeMaker treeMaker;

    /**
     * 提供了创建标识符的方法
     */
    private Names names;

    @Override
    public SourceVersion getSupportedSourceVersion() {
        return SourceVersion.latestSupported();
    }

    @Override
    public synchronized void init(ProcessingEnvironment processingEnv) {
        super.init(processingEnv);
        this.trees = JavacTrees.instance(processingEnv);
        Context context = ((JavacProcessingEnvironment) processingEnv).getContext();
        this.treeMaker = TreeMaker.instance(context);
        this.names = Names.instance(context);
    }

    @Override
    public boolean process(Set<? extends TypeElement> annotations, RoundEnvironment roundEnv) {
        // 获取被@Getter注解标记的所有元素(这个元素可能是类、变量、方法等等)
        Set<? extends Element> set = roundEnv.getElementsAnnotatedWith(Describe.class);
        for (Element element : set) {
            if (element instanceof Symbol.VarSymbol) {
                Symbol symbol = ((Symbol.VarSymbol) element).owner;
                for (Tree tree : trees.getPath(symbol)) {
                    if (tree instanceof JCTree.JCCompilationUnit) {
                        JCTree.JCCompilationUnit jccu = (JCTree.JCCompilationUnit) tree;
                        JCTree.JCImport jcImport = treeMaker.Import(treeMaker.Select(
                                treeMaker.Ident(getNameFromString("com.spring.dict.maven")), getNameFromString("DictUtil")), false);
                        jccu.defs = jccu.defs.prepend(jcImport);

                        jccu.accept(new TreeTranslator() {
                            @Override
                            public void visitClassDef(JCTree.JCClassDecl jcClassDecl) {
                                Describe describe = element.getAnnotation(Describe.class);
                                JCTree.JCMethodDecl jcMethodDecl = generateCodeDescribe(describe.args(), ((Symbol.VarSymbol) element).name);
                                jcClassDecl.defs = jcClassDecl.defs.prepend(jcMethodDecl);
                                super.visitClassDef(jcClassDecl);
                            }
                        });

                    }
                }
            }
        }
        return true;
    }

    JCTree.JCMethodDecl generateCodeDescribe(String[] args, Name name) {

        String funName = "get" + getMethodName(name.toString()) + "Describe";
        ListBuffer<JCTree.JCStatement> jcStatements = new ListBuffer<>();

        ListBuffer<JCTree.JCExpression> fieldNames = new ListBuffer<>();
        for (String arg : args) {
            fieldNames.append(treeMaker.Literal(arg));
        }


        JCTree.JCExpressionStatement printLiteral = treeMaker.Exec(treeMaker.Apply(
                List.of(memberAccess("java.lang.String")),
                memberAccess("DictUtil.getDict"),
                List.of(treeMaker.NewArray(memberAccess("java.lang.String"), List.nil(), fieldNames.toList()), memberAccess("this." + name))
                )
        );

        jcStatements.append(treeMaker.Return(printLiteral.getExpression()));
        //设置方法体
        JCTree.JCBlock jcBlock = treeMaker.Block(0, jcStatements.toList());

        JCTree.JCMethodDecl jcMethodDecl = treeMaker.MethodDef(
                treeMaker.Modifiers(Flags.PUBLIC),
                getNameFromString(funName),
                memberAccess("java.lang.String"),
                List.nil(),
                List.nil(),
                List.nil(),
                jcBlock,
                null
        );
        return jcMethodDecl;
    }

    /**
     * 首字母大写
     *
     * @param str
     * @return
     */
    private static String getMethodName(String str) {
        String[] cs = str.split("");
        cs[0] = cs[0].toUpperCase();
        return String.join("", cs);
    }


    private JCTree.JCExpression memberAccess(String components) {
        String[] componentArray = components.split("\\.");
        JCTree.JCExpression expr = treeMaker.Ident(getNameFromString(componentArray[0]));
        for (int i = 1; i < componentArray.length; i++) {
            expr = treeMaker.Select(expr, getNameFromString(componentArray[i]));
        }
        return expr;
    }

    private Name getNameFromString(String s) {
        return names.fromString(s);
    }
}
